#include "mem_pool.hh"
#include "ringbuffer.hh"

#include <iostream>

#define MEMORY_POOL_ENABLE 0

namespace kratos {
namespace network {

Trunk::Trunk(std::size_t elementSize)
    : count_(0), elementSize_(elementSize), totalSize_(0) {
    memList_.resize(64);
}

Trunk::~Trunk() {
    for (std::size_t i = 0; i < count_; i++) {
        free(memList_[i]);
    }
}

char* Trunk::allocate() {
    std::lock_guard<std::mutex> lock(mutex_);
    if (count_) {
        auto index = count_ - 1;
        auto pointer = memList_[index];
        memList_[index] = nullptr;
        count_ = index;
        return pointer;
    } else {
        auto pointer = reinterpret_cast<char*>(malloc(elementSize_));
        if (!pointer) {
            return nullptr;
        }
        totalSize_.fetch_add(elementSize_);
        return pointer;
    }
}

void Trunk::dispose(char* pointer) {
    std::lock_guard<std::mutex> lock(mutex_);
    if (count_ >= memList_.size()) {
        memList_.resize(memList_.size() * 2);
    }
    memList_[count_] = pointer;
    count_ += 1;
}

std::size_t Trunk::getTotalHeapSize() {
    return totalSize_;
}

MemoryPool::MemoryPool() {
    for (std::size_t size = 16; size <= MAX_ELEMENT_SIZE; size <<= 1) {
        trunkList_.push_back(new Trunk(size + sizeof(std::uint32_t)));
    }
}

MemoryPool::~MemoryPool() {
    for (auto trunk : trunkList_) {
        delete trunk;
    }
    trunkList_.clear();
}

#if MEMORY_POOL_ENABLE
static std::size_t log2(std::size_t size) {
    std::size_t n = 0;
    while (size > 1) {
        size >>= 1;
        n++;
    }
    return n;
}
#endif // MEMORY_POOL_ENABLE

char* MemoryPool::allocate(std::size_t size) {
#if !MEMORY_POOL_ENABLE
    return reinterpret_cast<char*>(malloc(size));
#else
    if (size < 16) {
        size = 16;
    }
    if (size > MAX_ELEMENT_SIZE) {
        auto pointer = reinterpret_cast<std::uint32_t*>(malloc(size + sizeof(std::uint32_t)));
        if (!pointer) {
            return nullptr;
        }
        *pointer = 0;
        return reinterpret_cast<char*>(pointer) + sizeof(std::uint32_t);
    }
    size = roundpow2(size);
    auto index = log2(size) - 4;
    auto pointer = reinterpret_cast<std::uint32_t*>(trunkList_[index]->allocate());
    *pointer = index + 1;
    return reinterpret_cast<char*>(pointer) + sizeof(std::uint32_t);
#endif // MEMORY_POOL_ENABLE
}

void MemoryPool::dispose(char * pointer) {
#if !MEMORY_POOL_ENABLE
    free(pointer);
    return;
#else
    if (!pointer) {
        return;
    }
    auto index = *reinterpret_cast<std::uint32_t*>(pointer - sizeof(std::uint32_t));
    if (!index) {
        free(pointer - sizeof(std::uint32_t));
        return;
    } else {
        index -= 1;
    }
    trunkList_[index]->dispose(pointer - sizeof(std::uint32_t));
#endif // MEMORY_POOL_ENABLE
}

std::size_t MemoryPool::getTotalHeapSize() {
    std::size_t totalHeapSize = 0;
    for (auto* trunk : trunkList_) {
        totalHeapSize += trunk->getTotalHeapSize();
    }
    return totalHeapSize;
}

}
}
